"""传递函数的频率特性"""
# from math import pi
import numpy as np
import matplotlib.pyplot as plt
plt.rcParams['font.family'] = ['Times New Roman', 'SimHei', ]

def poly(a, b):
    """计算两个多项式的积的系数列表"""
    """a和b分别是两个多项式的系数列表"""
    product = np.zeros(a.size + b.size - 1, dtype=np.float64)
    index = 0
    for factor in a:
        product[index:(b.size + index)] += factor*b
        index += 1
    return product

class System():
    """用传递函数分子和分母，描述一个动力学系统"""
    
    def __init__(self, num, den, **Bode_option):
        '''
        num：分子多项式。从0：end列分别表示从高阶到低阶排列的分子多项式系数；
        den：分母多项式。从0：end列分别表示从高阶到低阶排列的分母多项式系数；
        '''
        self.num = np.array(num, dtype=np.float64)
        self.num.shape = (self.num.size,)
        self.den = np.array(den, dtype=np.float64)
        self.den.shape = (self.den.size,)
        self.order = max(self.num.size, self.den.size)
        self.Bode_option = {'ShowPhase':False, 'grid':True, 
            'yAxis':'dB', 'xAxis':'log', 'Title':'', 'xLabel':'Frequency/Hz', 
            'yLabel':'Amplitude/dB', 'PhaseyLabel':'Phase/deg', }
        if Bode_option:
            for k, v in Bode_option.items():
                self.Bode_option[k] = v

    def __mul__(self, other_system):
        """重载运算符*：self和other_system串联的传函"""
        cascade_num = poly(self.num, other_system.num)
        cascade_den = poly(self.den, other_system.den)
        return System(cascade_num, cascade_den)
    
    def __truediv__(self, other_system):
        """重载运算符/：self和other_system传函的相除"""
        new_num = poly(self.num, other_system.den)
        new_den = poly(self.den, other_system.num)
        return System(new_num, new_den)

    def frequency_charac(self, freq):
        """根据传函的分子和分母多项式系数计算指定频率处的值"""
        ''' freq：指定待计算的频率/Hz '''
        accumulator = freq*1j*2*np.pi
        omegas = np.ones(freq.size, dtype=np.complex128)
        gains = np.zeros( (2, freq.size), dtype=np.complex128)
        index = 0
        num_size = self.num.size
        den_size = self.den.size
        while True:
            if index < num_size:
                # 矩阵（二维数组）切片得到矢量（一维数组）
                gains[0, :] += self.num[num_size-index-1]*omegas# 逐点计算
            if index < den_size:
                # 矩阵一定要用两个索引来访问，哪怕是shape=(1,n)的矩阵
                # 当索引数量不够时，缺失的索引默认为完整切片(:)
                # 缺省索引的规则类似函数形参默认值：有默认值的参数只能放在后面
                gains[1, :] += self.den[den_size-index-1]*omegas
            index += 1
            if index == self.order:
                break
            omegas *= accumulator
        return gains[0, :]/gains[1, :]

    # bad practice: 'trans_funcs' is mutable!
    # def Bode(self, freq, trans_funcs=[]):
    def Bode(self, freq):
        """绘制系统的Bode图"""
        if len(freq) == 2:
            start = np.log10(freq[0])
            stop = np.log10(freq[1])
            num = int((stop - start)*20) + 1
            freq = np.logspace(start, stop, num, dtype=np.float64)
            temp = self.frequency_charac(freq)
        else:
            temp = self.frequency_charac(np.array(freq))
        
        if self.Bode_option['ShowPhase'] == True:
            # ShowPhase特性尚未完成
            fig, axs = plt.subplots(2, 1)# y轴共线
            ax_amplitude = axs[0]
            ax_phase = axs[1]
        else:
            fig, ax_amplitude = plt.subplots()

        # phase = np.arctan(temp.real/temp.imag)    
        # axs[0].plot(Freq, amplitude)
        # axs[0].grid(self.Bode_option['grid'])
        # axs[0].set(ylabel=self.Bode_option['yLabel'], fontsize=11)
        # axs[1].plot(Freq, phase)
        # axs[1].grid(self.Bode_option['grid'])
        # axs[1].set(xlabel=self.Bode_option['xLabel'], \
        #     ylabel=self.Bode_option['PhaseyLabel'], fontsize=11)

        amplitude = np.abs(temp)
        if self.Bode_option['yAxis'] == 'dB':
            amplitude = 20*np.log10(amplitude)
            if self.Bode_option['xAxis'].lower() == 'log':
                ax_amplitude.semilogx(freq, amplitude)
            else:
                ax_amplitude.plot(freq, amplitude)
        elif self.Bode_option['yAxis'] == 'log':
            self.Bode_option['yLabel'] = 'log(Amplitude)'
            if self.Bode_option['xAxis'].lower() == 'log':
                ax_amplitude.loglog(freq, amplitude)
            else:
                ax_amplitude.semilogy(freq, amplitude)
        else:# 线性 y 轴
            self.Bode_option['yLabel'] = 'Amplitude'
            if self.Bode_option['xAxis'].lower() == 'log':
                ax_amplitude.semilogx(freq, amplitude)
            else:
                ax_amplitude.plot(freq, amplitude)
        ax_amplitude.set_xlabel(self.Bode_option['xLabel'], fontsize=11)
        ax_amplitude.set_ylabel(self.Bode_option['yLabel'], fontsize=11)
        ax_amplitude.grid(self.Bode_option['grid'])
        # plt.tick_params(axis='both', labelsize=11)
        # plt.show()

if __name__=='__main__':
    # tao1,tao2 = 2*np.pi*10, 2*np.pi*500
    # a = np.array( [1, tao1], dtype=np.float64)
    # b = np.array( [1, tao2], dtype=np.float64)
    # product = poly(a, b)
    # print(product)
    # low_pass_filter = System([tao1*tao2], product, xAxis='linear')
    # print(low_pass_filter.num.size)
    # print(low_pass_filter.den.size)
    # low_pass_filter.Bode(list(range(1000)))
    tao1, tao2 = 2*np.pi*5e4, 2*np.pi*1.5e5
    bp_filter = System([tao2, 0],\
        [1, tao1+tao2, tao1*tao2], yAxis='')
    bp_filter.Bode([1e3, 5e5])
    plt.show()
